package timer

import (
	"container/heap"
	"fmt"
	"sync"
	"time"
)

type TimerTask struct {
	task     func()
	interval time.Duration
	expire   int64
	stop     bool
}

func NewTimerTask(task func(), interval time.Duration) *TimerTask {
	return &TimerTask{task: task, interval: interval}
}

type TimerTaskQueue []*TimerTask

func (q TimerTaskQueue) Len() int            { return len(q) }
func (q TimerTaskQueue) Less(i, j int) bool  { return q[i].expire < q[j].expire }
func (q TimerTaskQueue) Swap(i, j int)       { q[i], q[j] = q[j], q[i] }
func (q TimerTaskQueue) Top() *TimerTask     { return q[0] }
func (q *TimerTaskQueue) Pop() interface{}   { r := (*q)[len(*q)-1]; *q = (*q)[0 : len(*q)-1]; return r }
func (q *TimerTaskQueue) Push(x interface{}) { *q = append(*q, x.(*TimerTask)) }

func (q *TimerTaskQueue) Dump() {
	for _, val := range *q {
		fmt.Println(val)
	}
}

type Timer struct {
	index     int32
	mtx       sync.Mutex
	tick      *time.Ticker
	taskIndex map[int32]*TimerTask
	taskQueue TimerTaskQueue
	done      chan struct{}
}

func NewTimer(interval int) *Timer {
	return &Timer{
		tick:      time.NewTicker(time.Duration(interval) * time.Millisecond),
		taskIndex: make(map[int32]*TimerTask),
		taskQueue: make(TimerTaskQueue, 0),
		done:      make(chan struct{}),
	}
}

func (t *Timer) Add(task *TimerTask) int32 {
	task.expire = time.Now().UnixNano() + task.interval.Nanoseconds()
	t.mtx.Lock()
	heap.Push(&t.taskQueue, task)
	t.index++
	t.taskIndex[t.index] = task
	t.mtx.Unlock()
	return t.index
}

func (t *Timer) Del(id int32) {
	t.mtx.Lock()
	task, ok := t.taskIndex[id]
	if ok {
		task.stop = true
		delete(t.taskIndex, id)
	}
	t.mtx.Unlock()
}

func (t *Timer) Elapse() {
	now := time.Now().UnixNano()
	size := t.taskQueue.Len()
	for size > 0 && t.taskQueue.Top().expire < now {
		top := heap.Pop(&t.taskQueue).(*TimerTask)
		if !top.stop {
			top.task()
			top.expire = top.interval.Nanoseconds() + now
			heap.Push(&t.taskQueue, top)
		}
		size = t.taskQueue.Len()
	}
}

func (t *Timer) Start() {
	go func() {
		defer t.tick.Stop()
		for {
			select {
			case <-t.tick.C:
				t.Elapse()
			case <-t.done:
				return
			}
		}
	}()
}

func (t *Timer) Stop() {
	t.done <- struct{}{}
}
